"""
Query understanding node - parses user's natural language query.
"""
from typing import Dict, Any
from langchain_core.prompts import ChatPromptTemplate
from datetime import datetime, timedelta
import json
from ..tools.llm_factory import create_llm


def understand_query(state: Dict[str, Any]) -> Dict[str, Any]:
    """
    Parse and understand the user's query.

    Extracts:
    - Target metric
    - Time range
    - Baseline period
    - Dimensions to analyze

    Args:
        state: Current state dict

    Returns:
        Updated state dict
    """
    query = state.get('query', '')
    config = state.get('config', {})

    # Create LLM using factory (supports OpenRouter and OpenAI)
    llm = create_llm(config=config)

    # Create prompt
    prompt = ChatPromptTemplate.from_messages([
        ("system", """You are an expert data analyst. Parse the user's query to extract:
1. The target metric to analyze
2. The time range for analysis (current period)
3. The baseline period for comparison
4. Any specific dimensions mentioned

Today's date is {today}.

Return a JSON object with these fields:
{{
    "target_metric": "metric name",
    "current_start": "YYYY-MM-DD",
    "current_end": "YYYY-MM-DD",
    "baseline_start": "YYYY-MM-DD",
    "baseline_end": "YYYY-MM-DD",
    "dimensions": ["dim1", "dim2"],
    "analysis_type": "anomaly_detection" or "trend_analysis" or "comparison"
}}

If dates are not specified, use reasonable defaults:
- Current period: last 7 days
- Baseline period: 7 days before current period
"""),
        ("user", "{query}")
    ])

    # Invoke LLM
    today = datetime.now().strftime('%Y-%m-%d')
    chain = prompt | llm
    response = chain.invoke({"query": query, "today": today})

    # Parse response
    try:
        result = json.loads(response.content)
    except json.JSONDecodeError:
        # Fallback to defaults
        current_end = datetime.now()
        current_start = current_end - timedelta(days=7)
        baseline_end = current_start - timedelta(days=1)
        baseline_start = baseline_end - timedelta(days=7)

        result = {
            "target_metric": "revenue",
            "current_start": current_start.strftime('%Y-%m-%d'),
            "current_end": current_end.strftime('%Y-%m-%d'),
            "baseline_start": baseline_start.strftime('%Y-%m-%d'),
            "baseline_end": baseline_end.strftime('%Y-%m-%d'),
            "dimensions": [],
            "analysis_type": "anomaly_detection"
        }

    # Update state
    state['target_metric'] = result.get('target_metric', 'revenue')
    state['time_range'] = (
        datetime.strptime(result['current_start'], '%Y-%m-%d'),
        datetime.strptime(result['current_end'], '%Y-%m-%d')
    )
    state['baseline_period'] = (
        datetime.strptime(result['baseline_start'], '%Y-%m-%d'),
        datetime.strptime(result['baseline_end'], '%Y-%m-%d')
    )

    # Store dimensions if specified, otherwise will be auto-detected from data
    if result.get('dimensions'):
        state['dimensions'] = result['dimensions']

    state['next_action'] = 'fetch_data'

    # Store LLM response for debugging
    if 'llm_responses' not in state:
        state['llm_responses'] = []
    state['llm_responses'].append(response.content)

    return state
